import numpy as np
from typing import List, Dict, Optional
import sys
import os

# Add the src directory to the path so we can import our modules
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from vocab import Vocab
from embeddings import Embeddings
from transformer import TransformerBlock
from output_projection import OutputProjection
from llm import LLM

# Constants - Reverted to original, smaller values to prevent overfitting
MAX_SEQ_LEN = 80
EMBEDDING_DIM = 128
HIDDEN_DIM = 256
NUM_HEADS = 4

def main():
    # Mock input - test conversational format
    string = "用户: 山是如何形成的?"

    # Extract all unique characters from training data to create vocabulary
    vocab_set = set()
    
    # Add special tokens that should be treated as whole units
    vocab_set.add("</s>")
    vocab_set.add("用户:")
    vocab_set.add("助手:")
    
    # Pre-training data - simple text completion patterns in Chinese
    pretraining_data = [
        "太阳从东方升起，在西方落下 </s>",
        "水因为重力作用从高处流向低处 </s>",
        "鸟类用翅膀在空中飞翔 </s>",
        "鱼类在河流、湖泊和海洋中游泳 </s>",
        "树木长得高大并长出叶子 </s>",
        "雨水从天空中的云层降落 </s>",
        "火是热的并能产生光亮 </s>",
        "冰是冻结的水，加热后会融化 </s>",
        "山脉是高耸的岩石构造 </s>",
        "月球围绕地球运转 </s>",
    ]

    chat_training_data = [
        # Conversational instruction-following data in Chinese
        "用户: 为什么会下雨? 助手: 雨是由云中的水蒸气凝结成水滴，当水滴太重无法悬浮在空中时降落形成的 </s>",
        "用户: 山是如何形成的? 助手: 山脉主要是通过地质构造运动或火山活动在漫长的地质时期中形成的 </s>",
        "用户: 什么是亚马逊雨林? 助手: 亚马逊雨林是地球上生物多样性最丰富的地区之一，拥有无数的物种 </s>",
        "用户: 水在什么温度下沸腾? 助手: 在标准大气压下，水在100摄氏度时沸腾 </s>",
        "用户: 月球绕地球运行需要多长时间? 助手: 月球大约每27.3天绕地球运行一周 </s>",
        "用户: 你好! 助手: 你好! 我能帮你什么吗? </s>",
        "用户: 谢谢。 助手: 不客气! 我很乐意帮忙 </s>",
    ]
    
    # Corrected tokenization: Process all training examples for vocabulary (character-level)
    all_data = pretraining_data + chat_training_data
    for text in all_data:
        # Split by space to correctly handle special tokens like "用户:" and "</s>"
        parts = text.split(' ')
        for part in parts:
            if part in vocab_set:
                continue # Special tokens are already in the set
            # For regular text, add each character to the vocab
            for char in part:
                vocab_set.add(char)

    vocab_words = sorted(list(vocab_set))
    vocab = Vocab(vocab_words)

    # Initialize the LLM with the corrected components, smaller architecture, and correct tokenizer mode
    llm = LLM(vocab, [
        Embeddings(vocab, EMBEDDING_DIM),
        TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, NUM_HEADS),
        TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, NUM_HEADS),
        TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, NUM_HEADS),
        OutputProjection(EMBEDDING_DIM, len(vocab.words)),
    ], tokenizer_mode='char') # Specify character mode for Chinese

    print("\n=== 模型信息 ===")
    print(f"网络架构: {llm.network_description()}")
    print(f"词汇表大小: {len(vocab.words)}")
    
    print("\n=== 训练前 ===")
    print(f"输入: {string}")
    print(f"输出: {llm.predict(string)}")
    
    print("\n=== 预训练模型 ===")
    print(f"在 {len(pretraining_data)} 个样本上预训练 100 轮，学习率 0.0005")
    llm.train(pretraining_data, 100, 0.0005)
    
    print("\n=== 指令微调 ===")
    print(f"在 {len(chat_training_data)} 个样本上进行指令微调 100 轮，学习率 0.0001")
    llm.train(chat_training_data, 100, 0.0001)
    
    print("\n=== 训练后 ===")
    print(f"输入: {string}")
    result = llm.predict(string)
    print(f"输出: {result}")
    print("======================\n")

    # Interactive mode for user input
    print("\n--- 交互模式 ---")
    print("输入提示词并按回车键生成文本。")
    print("输入 'exit' 退出程序。")
    
    while True:
        try:
            user_input = input("\n请输入: ").strip()
            
            if user_input.lower() == "exit":
                print("退出交互模式。")
                break
            
            formatted_input = f"用户: {user_input}"
            prediction = llm.predict(formatted_input)
            print(f"模型输出: {prediction}")
        except KeyboardInterrupt:
            print("\n退出交互模式。")
            break
        except EOFError:
            print("\n退出交互模式。")
            break

if __name__ == "__main__":
    main()
